import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from sklearn.model_selection import train_test_split

# 创建虚拟时间序列数据集
# 假设我们有一个包含788个时间步的时间序列，每个时间步有60个特征
num_samples = 1000
sequence_length = 788
input_features = 60

# 1. 准备数据
file_name = 'feature/JDT.csv'
# 读取csv
data = pd.read_csv(file_name, header=None)
data = pd.DataFrame(data)
data = data.drop([0])
# 继续读取csv，并且添加到data中
# temp = pd.read_csv('feature/ML.csv', header=None)
# temp = pd.DataFrame(temp)
# temp = temp.drop([0])
# data = pd.concat([data, temp], axis=0, ignore_index=True)
# temp = pd.read_csv('feature/PDE.csv', header=None)
# temp = pd.DataFrame(temp)
# temp = temp.drop([0])
# data = pd.concat([data, temp], axis=0, ignore_index=True)
# temp = pd.read_csv('feature/LC.csv', header=None)
# temp = pd.DataFrame(temp)
# temp = temp.drop([0])
# data = pd.concat([data, temp], axis=0, ignore_index=True)
# temp = pd.read_csv('feature/EQ.csv', header=None)
# temp = pd.DataFrame(temp)
# temp = temp.drop([0])
# data = pd.concat([data, temp], axis=0, ignore_index=True)
datasets = data.iloc[:, :-1]
labels = data.iloc[:, -1]

# 将特征数据转为数组
datasets = np.array(datasets)
# print(datasets)
# 标签的转换为0、1
labels = np.array(labels)

for i in range(len(labels)):
    if labels[i] == "b'clean'":
        labels[i] = np.int32(0)
    else:
        labels[i] = np.int32(1)

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(datasets, labels, test_size=0.2, random_state=42)

# x转置
X_train = X_train.T
X_test = X_test.T


# 创建CNN模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv1d(in_channels=input_features, out_channels=32,
                      kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
        )
        self.max = nn.MaxPool1d(kernel_size=2)
        self.fc1 = nn.Linear(16, 64)  # 根据输入数据形状调整这里的输入维度
        self.fc2 = nn.Linear(64, 2)  # 2表示二分类问题

    def forward(self, x):
        x = self.layer1(x)
        x = x.T  # 根据输入数据形状调整这里的输入维度
        x = self.max(x)
        x = self.fc1(x)
        x = self.fc2(x)
        return x


# 实例化模型、损失函数和优化器
model = SimpleCNN()
criterion = nn.CrossEntropyLoss()

optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    optimizer.zero_grad()
    inputs = torch.tensor(X_train.astype(float))
    labels = torch.tensor(y_train.astype(int))
    outputs = model(inputs.float())
    print(outputs.shape,labels.shape)
    loss = criterion(outputs, labels.long())
    loss.backward()
    optimizer.step()
    print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item()}")

# 在测试集上评估模型
model.eval()
with torch.no_grad():
    test_inputs = torch.tensor(X_test.astype(float))
    test_labels = torch.tensor(y_test.astype(int))
    test_outputs = model(test_inputs.float())
    predicted = torch.argmax(test_outputs, dim=1)
    accuracy = torch.sum(predicted == test_labels).item() / len(test_labels)
    print(f"Test Accuracy: {accuracy}")
